package ffte.core.flat

import spinal.core._
import spinal.lib._

import ffte.types.{FixReal,FixComplex}
import ffte.algorithm.Prime
import ffte.property.{getTwiddleResolution}

/**
 * Twiddle factor trait for FFT operations
 *
 * This trait provides a unified interface for different types of twiddle factors
 * used in FFT computations. Twiddle factors are complex roots of unity that
 * appear in FFT algorithms, particularly in Cooley-Tukey based decompositions.
 *
 * Twiddle Factor Types:
 * - Simple factors: ±1, ±j (no multiplication required)
 * - Complex factors: General complex exponentials (require multiplication)
 * - Optimized factors: Special cases with reduced multiplication count
 *
 * Hardware Optimization:
 * - Simple factors avoid expensive complex multiplications
 * - Special cases like 8-point DFT use optimized algorithms
 * - Fixed-point precision managed through configurable resolution
 * - Multiplication count minimized for hardware efficiency
 *
 * Mathematical Background:
 * Twiddle factors are of the form W_N^k = e^(-j*2π*k/N) where:
 * - N: Transform size
 * - k: Twiddle factor index
 * - Negative exponent for forward FFT
 */
trait twiddle {
    /**
     * Get simplified twiddle factor representation
     *
     * Attempts to reduce complex twiddle factors to simpler forms
     * when possible. For TW factors, this checks if they can be
     * simplified to basic constants (±1, ±j) or reduced forms.
     *
     * @return Simplified twiddle factor (T1, T_1, TJ, T_J, or reduced TW)
     */
    def simple = this match {
        case w: TW => w._simple
        case _ => this
    }

    /**
     * Check if multiplication is needed
     *
     * Determines whether this twiddle factor requires complex multiplication
     * or if it can be implemented with simple operations (sign changes,
     * swapping real/imaginary parts).
     *
     * @return true if complex multiplication is required, false for simple factors
     */
    def needMul = this.simple match {
        case w: TW => true
        case _ => false
    }

    /**
     * Get multiplication cost
     *
     * Returns the number of real multiplications required for this
     * twiddle factor. Complex multiplication typically requires 4
     * real multiplications, but special cases like 8-point DFT
     * can be optimized to use only 2.
     *
     * @return Number of real multiplications required
     */
    def mulTimes = this.simple match {
        case TW(_, n) => if (n == 8) 2 else 4
        case _ => 0
    }

    /**
     * Get twiddle factor bit width
     *
     * Returns the bit width used for twiddle factor representation.
     * Uses reduced precision compared to data path to optimize
     * hardware resources while maintaining accuracy.
     *
     * @return Bit width for twiddle factor representation
     */
    def width = 1 - getTwiddleResolution()

    /**
     * Perform complex multiplication with twiddle factor
     *
     * Multiplies the input complex number by this twiddle factor
     * using optimized algorithms for different cases:
     *
     * Simple Cases:
     * - T1: Identity (no operation)
     * - T_1: Negation (sign change)
     * - TJ: Multiply by j (swap real/imaginary)
     * - T_J: Multiply by -j (swap and negate)
     *
     * Optimized 8-point Case:
     * Uses the identity: W_8^1 = (1/√2)(1 - j)
     * Requires only 2 real multiplications instead of 4
     *
     * General Case:
     * Standard complex multiplication with cos/sin values
     *
     * @param d Input complex number to multiply
     * @return Result of multiplication with appropriate precision management
     */
    def mul(d: FixComplex) = this.simple match {
        // Simple twiddle factors requiring no multiplication
        case T1 => d      // Identity: W_N^0 = 1
        case T_1 => -d    // Negation: W_N^(N/2) = -1
        case TJ => d.j    // Multiply by j: W_N^(N/4) = j
        case T_J => -d.j  // Multiply by -j: W_N^(3N/4) = -j

        // Complex twiddle factors requiring multiplication
        case TW(k, n) => {
            val r = 1 << (width - 1)  // Scaling factor for fixed-point representation

            if (n == 8) {
                /**
                 * Optimized 8-point DFT twiddle factors
                 *
                 * Uses the mathematical identity:
                 * W_8^1 = e^(-j*π/4) = (1/√2)(1 - j)
                 * W_8^3 = e^(-j*3π/4) = (1/√2)(-1 - j)
                 * W_8^5 = e^(-j*5π/4) = (1/√2)(-1 + j)
                 * W_8^7 = e^(-j*7π/4) = (1/√2)(1 + j)
                 *
                 * This optimization reduces 4 real multiplications to 2
                 * by using the structure of the cosine/sine values.
                 */
                val cospi_4 = Math.cos(Math.PI / 4)
                val f = FixReal(width).fromInt(Math.round(cospi_4 * r).toInt)

                // Temporary complex values for optimized computation
                val c = FixComplex(d.width + 1)
                val e = FixComplex(d.width)

                // Compute (d.re - d.im) + j(d.re + d.im) = d * (1 + j)
                c.re := d.re.expand - d.im
                c.im := d.re.expand + d.im

                // Multiply by cos(π/4) = sin(π/4) = 1/√2
                e := (c * f).cut(d.width, width - 1, d.resolution)

                // Map k value to appropriate twiddle factor
                k match {
                    case 1 => e      // W_8^1 = (1/√2)(1 - j)
                    case 3 => e.j    // W_8^3 = (1/√2)(-1 - j) = j * W_8^1
                    case 5 => -e     // W_8^5 = (1/√2)(-1 + j) = -W_8^1
                    case 7 => -e.j   // W_8^7 = (1/√2)(1 + j) = -j * W_8^1
                    case _ => FixComplex(d.width).zero  // Invalid case, return zero
                }
            } else {
                /**
                 * General twiddle factor implementation
                 *
                 * Uses standard complex multiplication: d * (cos + j*sin)
                 * where cos = cos(2π*k/n) and sin = -sin(2π*k/n)
                 *
                 * The twiddle factors are pre-computed and stored as
                 * fixed-point values to optimize hardware implementation.
                 */
                val c = FixComplex(width)
                val cos = Math.cos(2 * k * Math.PI / n)
                val sin = Math.sin(2 * k * Math.PI / n)

                // Load fixed-point cosine and sine values
                c.re.d.raw := S(Math.round(cos * r).toInt, width bits)
                c.im.d.raw := S(Math.round(sin * r).toInt, width bits)

                // Perform complex multiplication with precision management
                (d * c).cut(d.width, width - 1, d.resolution)
            }
        }
    }
    /**
     * Registered rotation operation
     *
     * Performs twiddle factor multiplication with registered output.
     * This breaks combinatorial paths and provides proper timing
     * control for pipelined FFT implementations.
     *
     * The operation is only performed when the clock enable signal
     * is asserted, allowing precise control over computation timing.
     *
     * @param d Input complex number to rotate
     * @param ce Clock enable signal controlling when to perform rotation
     * @return Registered output of the rotation operation
     */
    def Rotate(d: FixComplex, ce: Bool) = {
        val ra = new Area {
            /**
             * Output register for rotated value
             *
             * Stores the result of twiddle factor multiplication.
             * Initialized to zero to avoid undefined values.
             */
            val r = RegInit(FixComplex(d.width).zero)

            /**
             * Perform rotation when clock enable is asserted
             *
             * The multiplication is only performed when ce is high,
             * providing synchronous control over the operation.
             */
            when(ce) {
                r := mul(d)
            }
        }
        ra.r
    }
}

/**
 * Twiddle factor T_1 (-1)
 *
 * Represents the twiddle factor W_N^(N/2) = -1.
 * Used for sign negation without requiring multiplication.
 */
case object T_1 extends twiddle

/**
 * Twiddle factor T1 (1)
 *
 * Represents the twiddle factor W_N^0 = 1.
 * Used for identity operation (no change to input).
 */
case object T1 extends twiddle

/**
 * Twiddle factor TJ (j)
 *
 * Represents the twiddle factor W_N^(N/4) = j.
 * Used for 90-degree rotation (swap real and imaginary parts).
 */
case object TJ extends twiddle

/**
 * Twiddle factor T_J (-j)
 *
 * Represents the twiddle factor W_N^(3N/4) = -j.
 * Used for 270-degree rotation (swap and change sign).
 */
case object T_J extends twiddle

/**
 * Complex twiddle factor TW(k,n)
 *
 * Represents a general twiddle factor W_n^k = e^(-j*2π*k/n) for arbitrary
 * transform size n and index k. This class handles the full range of
 * complex twiddle factors with optimization for special cases.
 *
 * Mathematical Representation:
 * TW(k,n) = cos(2π*k/n) - j*sin(2π*k/n)
 *
 * Optimization Strategy:
 * - Simplify to basic factors (±1, ±j) when possible
 * - Reduce by GCD to minimize computation
 * - Use special algorithms for n=8 case
 * - General complex multiplication for other cases
 *
 * @param k Twiddle factor index (0 ≤ k < n)
 * @param n Transform size (must be > 0)
 */
case class TW(k: Int, n: Int) extends twiddle {
    /**
     * Simplify twiddle factor to basic form
     *
     * Attempts to reduce the twiddle factor to simpler forms
     * that avoid expensive multiplication operations:
     *
     * Simplification Rules:
     * - W_n^0 = 1 → T1 (identity)
     * - W_n^(n/2) = -1 → T_1 (negation)
     * - W_n^(n/4) = j → TJ (90° rotation)
     * - W_n^(3n/4) = -j → T_J (270° rotation)
     * - Reduce by GCD for other cases
     *
     * The modulo operation (k+2n)%n ensures positive indices
     * and handles edge cases correctly.
     *
     * @return Simplified twiddle factor (T1, T_1, TJ, T_J, or reduced TW)
     */
    def _simple = {
        val vk = (k + 2 * n) % n  // Normalize to positive index

        vk match {
            case 0 => T1                           // W_n^0 = 1
            case vk if (vk * 2 == n) => T_1        // W_n^(n/2) = -1
            case vk if (vk * 4 == n) => TJ         // W_n^(n/4) = j
            case vk if (vk * 4 == 3 * n) => T_J    // W_n^(3n/4) = -j
            case _ => {
                // Reduce by GCD to minimize computation
                val g = Prime.gcd(vk, n)
                TW(vk / g, n / g)
            }
        }
    }
}


